//-----------------------------------------------------------------------------
//
// @file	Image.cpp
// @brief	Image include files, variables and classes
// @author	Vinod Melapudi
// @date	20Feb2010
//
//-----------------------------------------------------------------------------

#include "Image.h"

#include "OpenGLRender.h"

#include <png.h>


Image::Image()
{
	m_width = 0;
	m_height = 0;
	m_color_mode = GL_RGBA;
	m_data = NULL;
}

Image::~Image()
{
	if ( m_data )
		free(m_data);
}

Image::Image( const char * fileName, IMAGE_TYPE type )
{
	m_width = 0;
	m_height = 0;
	m_color_mode = GL_RGBA;
	m_data = NULL;

	Read( fileName, type );
}

bool Image::Read( const char* fileName, IMAGE_TYPE type )
{
	if ( m_data )
		free( m_data );

	switch( type )
	{
	case IMAGE_BMP:	return ReadBMP( fileName );	break;
	case IMAGE_PNG: return ReadPNG( fileName );	break;
	}

	fprintf( stderr, "Image type %d not handled", type );
	return false;
}

bool Image::Write( const char* fileName, IMAGE_TYPE type )
{
	switch( type )
	{
	case IMAGE_BMP:	return WriteBMP( fileName );	break;
	case IMAGE_PNG: return WritePNG( fileName );	break;
	}

	fprintf( stderr, "Image type %d not handled", type );
	return false;
}

//
// @brief	Read a BMP file to texture data
//			Mallocs char*, this should be free later
//
bool Image::ReadBMP( const char* fileName )
{
	TCHAR NPath[MAX_PATH];
	GetCurrentDirectory(MAX_PATH, NPath);

	// Read texture file
	FILE * file = fopen( fileName, "rb" );
	if ( file == NULL ) 
	{
		fprintf( stderr, "Can't open file '%s' for reading\n", fileName );
		return false;
	}

	// Skip the BMP headers (Assuming there is no junk bytes)
	BITMAPFILEHEADER bmFileHdr;
	BITMAPINFO bmInfo;
	fread((char *)&bmFileHdr, sizeof(BITMAPFILEHEADER), 1, file);
	fread((char *)&bmInfo, sizeof(BITMAPINFO), 1, file);
	fseek(file, bmFileHdr.bfOffBits, SEEK_SET);

	//	if ( bmInfo.bmiHeader.biWidth != bmInfo.bmiHeader.biHeight )
	//	{
	//		fprintf(stderr, "Texture should have same width & height for '%s'\n", fileName );
	//		return NULL;
	//	}

	m_width = bmInfo.bmiHeader.biWidth;
	m_height = bmInfo.bmiHeader.biHeight;
	m_color_mode = GL_BGR;

	// Allocate memory
	if ( m_data )
		free( m_data );
	long len = bmInfo.bmiHeader.biWidth * bmInfo.bmiHeader.biHeight * 3;
	m_data = (unsigned char*) malloc(len * sizeof(char));
	if ( m_data == NULL )
	{
		fprintf(stderr, "Can't allocate space for '%s'\n", fileName );
		fclose( file );
		return false;
	}

	long read = fread( m_data, sizeof(char), len, file );
	fclose( file );

	m_data[read-1] = '\0';
	return true;
}

// Write BGR data to BMP file
bool Image::WriteBMP( const char* fileName )
{
	// Make sure width of data is multiples of 4 (this is a requirement of BMP)
	if ( (m_width % 4) != 0 )
	{
		fprintf( stderr, "WriteBMP: The width of data to write is not multiple of 4." );
		return false;
	}


	// Write file
	FILE * file = fopen( fileName, "wb" );
	if ( file == NULL ) 
	{
		fprintf( stderr, "WriteBMP: Can't open file '%s' for writing\n", fileName );
		return false;
	}


	// Write BMP headers
	{
		size_t hs = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFO);

		BITMAPFILEHEADER bmFileHdr;
		bmFileHdr.bfType				= 0x4d42;
		bmFileHdr.bfSize				= hs + 3 * m_width * m_height;
		bmFileHdr.bfReserved1			= 0;
		bmFileHdr.bfReserved2			= 0;
		bmFileHdr.bfOffBits				= hs;

		BITMAPINFO bmInfo;
		bmInfo.bmiHeader.biSize			= 40;
		bmInfo.bmiHeader.biWidth		= m_width;
		bmInfo.bmiHeader.biHeight		= m_height;
		bmInfo.bmiHeader.biPlanes		= 1;
		bmInfo.bmiHeader.biBitCount		= 24;
		bmInfo.bmiHeader.biCompression	= 0;
		bmInfo.bmiHeader.biSizeImage	= 3 * m_width * m_height;
		bmInfo.bmiHeader.biXPelsPerMeter= 0;
		bmInfo.bmiHeader.biYPelsPerMeter= 0;
		bmInfo.bmiHeader.biClrUsed		= 0;
		bmInfo.bmiHeader.biClrImportant = 0;


		fwrite((char *)&bmFileHdr, sizeof(BITMAPFILEHEADER), 1, file);
		fwrite((char *)&bmInfo, sizeof(BITMAPINFO), 1, file);
	}

	// Write data into file
	fwrite(m_data, sizeof(char), 3*m_width*m_height, file);

	fclose( file );
}


// Load color to image
void Image::LoadColor( int w, int h, const std::vector<Color<short>>& colors )
{
	m_width = w;
	m_height = h;
	m_color_mode = GL_RGBA; 
	assert( colors.size() == m_width * m_height );

	if ( m_data )
		free(m_data);
	m_data = (unsigned char*) malloc( 4 * m_width * m_height );

	for( size_t i=0; i < colors.size(); i++ )
	{
		m_data[i*3+0] = (char)colors[i].r;
		m_data[i*3+1] = (char)colors[i].g;
		m_data[i*3+2] = (char)colors[i].b;
		m_data[i*3+3] = (char)colors[i].a;
	}
}

// Load viewport to image
bool Image::LoadViewport()
{
	// Write viewport to file
	GLint viewport[4];
	glGetIntegerv( GL_VIEWPORT, viewport );
	m_width = viewport[2];
	m_height = viewport[3];
	m_color_mode = GL_RGBA;
	int num_components = 4;

	if ( m_data ) 
		free( m_data );
	m_data = (unsigned char*) malloc( num_components * m_width * m_height );

	glReadBuffer(GL_BACK);
	glReadPixels( 0,0, m_width, m_height, GL_RGBA, GL_UNSIGNED_BYTE, m_data );

	// Flip data vertically
	int pixels_per_row = num_components * m_width;
	for( int i=0; i < floorf(m_height/2); i++ )
	{
		int k = m_height - i - 1;
		int iw = i*pixels_per_row;
		int kw = k*pixels_per_row;
		for( int j=0; j < pixels_per_row; j++ )
		{
			unsigned char c = m_data[iw+j];
			m_data[iw+j] = m_data[kw+j];
			m_data[kw+j] = c; 
		}
	}

	return true;
}

// Read data from PNG
bool Image::ReadPNG( const char* fileName )
{
	unsigned char sig[8];    // 8 is the maximum size that can be checked

	/* open file and test for it being a png */
	FILE *fp = fopen(fileName, "rb");
	if (!fp)
	{
		fprintf(stderr, "File %s could not be opened for reading", fileName);
		return false;
	}

	fread(sig, 1, 8, fp);
	if (!png_check_sig(sig, 8))
	{
		fprintf(stderr, "File %s is not recognized as a PNG file", fileName);
		fclose(fp);
		return false;
	}

	/* initialize stuff */
	png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

	if (!png_ptr)
	{
		fprintf(stderr, "png_create_read_struct failed");
		fclose(fp);
		return false;
	}

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)
	{
		fprintf(stderr, "png_create_info_struct failed");
		fclose(fp);
		return false;
	}

	if (setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr, "Error during init_io");
		fclose(fp);
		return false;
	}

	png_init_io(png_ptr, fp);
	png_set_sig_bytes(png_ptr, 8);

	png_read_info(png_ptr, info_ptr);

	m_width = png_get_image_width(png_ptr, info_ptr);
	m_height = png_get_image_height(png_ptr, info_ptr);
	png_byte color_type = png_get_color_type(png_ptr, info_ptr);
	png_byte bit_depth = png_get_bit_depth(png_ptr, info_ptr);

	int components = 0;
	assert( color_type == PNG_COLOR_TYPE_RGB || color_type == PNG_COLOR_TYPE_RGBA );
	if ( color_type == PNG_COLOR_TYPE_RGB )
	{
		components = 3;
		m_color_mode = GL_RGB;
	}
	else if ( color_type == PNG_COLOR_TYPE_RGBA )
	{
		components = 4;
		m_color_mode = GL_RGBA;
	}
	else
	{
		fprintf(stderr, "Color type %d not handled", color_type );
		fclose(fp);
		return false;
	}

	int number_of_passes = png_set_interlace_handling(png_ptr);
	png_read_update_info(png_ptr, info_ptr);


	/* read file */
	if (setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr, "Error during read_image");
		fclose(fp);
		return false;
	}

	assert( png_get_rowbytes(png_ptr, info_ptr) == m_width * components );

	// Set pointers for data and row majors
	png_byte* color_data = (png_byte*) malloc( sizeof(png_byte) * m_height * m_width * components );
	png_bytep* row_pointers = (png_bytep*) malloc(sizeof(png_bytep) * m_height);
	for (int y=0; y < m_height; y++)
		row_pointers[y] = &color_data[y * m_width * components];

	// Read image
	png_read_image(png_ptr, row_pointers);

	png_read_end(png_ptr, NULL);
	png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
	free( row_pointers );
	fclose(fp);

	if ( m_data )
		free(m_data);
	m_data = (unsigned char*) color_data;
}

// Write color data to PNG
bool Image::WritePNG( const char* fileName )
{
	// Get components size
	int components = 3;
	png_byte color_type = PNG_COLOR_TYPE_RGB;
	png_byte bit_depth = 8;

	if ( m_color_mode == GL_RGB )
	{
		components = 3;
		color_type = PNG_COLOR_TYPE_RGB;
	}
	else if ( m_color_mode == GL_RGBA )
	{
		components = 4;
		color_type = PNG_COLOR_TYPE_RGBA;
	}
	else
	{
		fprintf(stderr,"Color mode %d not handled", (int)m_color_mode);
		return false;
	}

	// open file create file 
	FILE *fp = fopen(fileName, "wb");
	if (!fp)
	{
		fprintf(stderr,"File %s could not be opened for writing", fileName);
		return false;
	}

	/* initialize stuff */
	png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (!png_ptr)
	{
		fprintf(stderr,"png_create_write_struct failed");
		fclose(fp);
		return false;
	}

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)		
	{
		fprintf(stderr,"png_create_info_struct failed");
		fclose(fp);
		return false;
	}

	if (setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"Error during init_io");
		fclose(fp);
		return false;
	}

	png_init_io(png_ptr, fp);


	// Write header 
	if (setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"Error during writing header");
		fclose(fp);
		return false;
	}

	png_set_IHDR(png_ptr, info_ptr, m_width, m_height,
		bit_depth, color_type, PNG_INTERLACE_NONE,
		PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

	png_write_info(png_ptr, info_ptr);


	// Write data
	if (setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"Error during writing bytes");
		fclose(fp);
		return false;
	}

	// Put data into row major arrays
	png_bytep* row_pointers = (png_bytep*) malloc(sizeof(png_bytep) * m_height);
	for (int y=0; y < m_height; y++)
		row_pointers[y] = (png_byte*) &m_data[y * m_width * components];

	// Write image
	png_write_image(png_ptr, row_pointers);


	// End write 
	if (setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"Error during end of write");
		fclose(fp);
		free( row_pointers );
		return false;
	}

	png_write_end(png_ptr, NULL);

	fclose(fp);
	free( row_pointers );

	return true;
}